"""
Single Shot MultiBox Detector
"""
import torch
import torch.nn as nn
import torch.nn.functional as F

from vortex.network.backbone.vggbase import VGGBase

from vortex.network.layers.anchor_creator import SSDAnchorCreator
from vortex.network.head.ssd import SSDDetect, SSDHead, MultiBoxLoss
from vortex.network.backbone import get_backbone_model

import vortex.data.transforms as VT
import vortex.data.transforms.image as VIT
from torchvision import transforms as FT


# -------------------------------------------------------------------------------
# NEW SSD Implementation
class SSD(nn.Module):
    def __init__(self, num_classes, image_size=(300, 300), backbone='vgg16', backbone_weights=None):
        super(SSD, self).__init__()
        num_classes_with_bg = num_classes + 1

        self.image_size = image_size

        self.backbone = get_backbone_model(backbone, backbone_weights)
        self.head = SSDHead(num_classes_with_bg)

        anchor_creator = SSDAnchorCreator()
        anchors = anchor_creator.create()

        self.detect_fn = SSDDetect(num_classes_with_bg, anchors)
        self.loss_fn = MultiBoxLoss(anchors)

    def get_data_transform(self, split):
        assert split in ['train', 'trainval', 'test', 'val']
        transforms = []

        if split in ['train', 'trainval']:
            transforms.append(VT.RandomPhotometric())
            transforms.append(VT.RandomExpand())
            transforms.append(VT.RandomScaledCrop())
            transforms.append(VT.RandomHorizontalFlip())
        
        transforms.append(VT.Resize(size=self.image_size))
        transforms.append(VT.ToTensor())
        transforms.append(VT.Normalize())

        transforms = VT.Compose(transforms)
        return transforms
    
    def get_inference_transform(self):
        """
        inference transforms take single image as input
        """
        resize = VIT.ImageResize(self.image_size)
        to_tensor = FT.ToTensor()
        normalize = FT.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        transforms = FT.Compose([resize, to_tensor, normalize])
        return transforms
    
    def detect(self, inputs):
        """
        detect objects on single preprocessed image
        """
        outputs = self.backbone(inputs)
        locs, confs = self.head(outputs)

        boxes, labels, scores = self.detect_fn.detect(locs, confs)

        # parse outputs
        boxes = boxes[0].detach().cpu().numpy()
        labels = labels[0].to('cpu').tolist()
        scores = scores[0].to('cpu').tolist()

        return boxes, labels, scores

    def forward(self, inputs, boxes, labels):
        """
        calculate loss for inputs
        """
        outputs = self.backbone(inputs)
        locs, confs = self.head(outputs)
        loss = self.loss_fn(locs, confs, boxes, labels)
        return loss
